
import os
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
import tensorflow as tf
from tensorflow import keras
from Model import makeDarkNet, makeDarkNet53, makeDarkNet2
from DataBase import getDataset, augmentationDataBase, groundTruthGenerator
from YoloLoss import yolo_loss
from tensorflow.keras import layers
import  tensorflow.keras.backend as K
import numpy as np
from tensorflow import keras

if __name__ == '__main__':

    startEpoch = 0
    bestLoss = 0.5
    # model = makeDarkNet53([1, 2, 3, 3, 2])
    model = makeDarkNet()
    # model = keras.models.load_model("BestModel_55epoch_0.06833954456262291loss.h5")
    # model.load_weights("savedModels\\BespiptModel_151epoch_0.05334455331787467loss.h5")
    # model.load_weights("savedModels\\darknet18_finaladd_2\\BestModel_139epoch_0.022783386027440428loss.h5")
    # model.load_weights("")
    # objNames = ('trash1', 'trash2')
    objNames = ("trash1", "trash2")
    # trainDataBase = getDataset('data\\trainDarkNet53', 'data\\trainDarkNet53Xml', objNames, 16)
    # valDataBase = getDataset("data\\valDarkNet53", "data\\valDarkNet53Xml", objNames, 2)
    # trainDataBase = getDataset('newData\\trainimg', 'newData\\trainlabel', objNames, 16)
    # valDataBase = getDataset("newData\\testimg", "newData\\testlabel", objNames, 1)

    trainDataBase = getDataset('data20230208\\trainimgs', 'data20230208\\trainlabels', objNames, 16)
    valDataBase = getDataset("data20230208\\testimgs", "data20230208\\testlabels", objNames, 1)

    trainDataBase = augmentationDataBase(trainDataBase)
    valDataBase = augmentationDataBase(valDataBase)
    trainGen = groundTruthGenerator(trainDataBase)
    valGen = groundTruthGenerator(valDataBase)

    # testGen = groundTruthGenerator(getDataset("data/test/*.jpg", "data/testxml/*.xml", 32))


    optimizer = None

    for epoch in range(startEpoch, 10000):
        if epoch >= 0:
            optimizer = keras.optimizers.Adam(learning_rate=1e-4, beta_1=0.9,
                                              beta_2=0.999, epsilon=1e-08)
        if epoch >= 100:
            optimizer = keras.optimizers.Adam(learning_rate=1e-5, beta_1=0.9,
                                              beta_2=0.999, epsilon=1e-08)
        if epoch >= 250:
            optimizer = keras.optimizers.Adam(learning_rate=1e-6, beta_1=0.9,
                                              beta_2=0.999, epsilon=1e-08)
        if epoch == 300:
            optimizer = keras.optimizers.Adam(learning_rate=1e-7, beta_1=0.9,
                                              beta_2=0.999, epsilon=1e-08)

        trainEpoch = 200
        trainTotalLoss = 0.
        for step in range(trainEpoch):
            img, detector_mask, matching_true_boxes, matching_classes_oh, true_boxes = next(trainGen)
            with tf.GradientTape() as tape:
                y_pred = model(img, training=True)
                loss, sub_loss = yolo_loss(detector_mask, \
                                           matching_true_boxes, matching_classes_oh, \
                                           true_boxes, y_pred)
            grads = tape.gradient(loss, model.trainable_variables)
            optimizer.apply_gradients(zip(grads, model.trainable_variables))

            # print(epoch, step, float(loss), float(sub_loss[0]), float(sub_loss[1]), float(sub_loss[2]))
            # if epoch > 30 and bestLoss > float(loss):
            #     # model.save("BestModel_{}epoch_{}loss.HDF5".format(epoch, float(loss)))
            #     model.save_weights("savedModels\\BestModel_{}epoch_{}loss.h5".format(epoch, float(loss)))
            #     bestLoss = (float(loss))
            trainTotalLoss += float(loss)
        trainLoss = trainTotalLoss / trainEpoch

        testEpoch = 100
        testTotalLoss = 0.
        for step in range(testEpoch):
            img, detector_mask, matching_true_boxes, matching_classes_oh, true_boxes = next(trainGen)
            y_pred = model(img, training=False)
            loss, sub_loss = yolo_loss(detector_mask, \
                                       matching_true_boxes, matching_classes_oh, \
                                       true_boxes, y_pred)
            # print(epoch, step, float(loss), float(sub_loss[0]), float(sub_loss[1]), float(sub_loss[2]))
            # if epoch > 30 and bestLoss > float(loss):
            #     # model.save("BestModel_{}epoch_{}loss.HDF5".format(epoch, float(loss)))
            #     model.save_weights("savedModels\\BestModel_{}epoch_{}loss.h5".format(epoch, float(loss)))
            #     bestLoss = (float(loss))
            testTotalLoss += float(loss)
        testLoss = testTotalLoss / testEpoch
        print(epoch, "train loss:", trainLoss, "test loss:", testLoss)

        testLoss = trainLoss
        if epoch > 0 and bestLoss > testLoss:
            model.save_weights("savedModels\\BestModel_{}epoch_{}loss.h5".format(epoch, testLoss))
            bestLoss = testLoss